#include <atomic>
#include <chrono>
#include <cstdint>
#include <thread>
#include <vector>
#include <iostream>
#include <stdexcept>

class MultiSlotSlidingWindowLimiter {
public:
    using clock = std::chrono::steady_clock;

    struct Slot {
        std::atomic<uint64_t> count;
        std::atomic<uint64_t> start_ms;
        Slot() : count(0), start_ms(0) {}
    };

    MultiSlotSlidingWindowLimiter(uint64_t window_size_ms, uint64_t capacity, size_t slot_count = 10)
        : window_size_ms_(window_size_ms), capacity_(capacity), slot_count_(slot_count), slots_(slot_count) 
    {
        if (window_size_ms == 0 || capacity == 0 || slot_count == 0)
            throw std::invalid_argument("window_size, capacity, slot_count must be > 0");
        slot_size_ms_ = window_size_ms_ / slot_count_;
        if (slot_size_ms_ == 0) slot_size_ms_ = 1; // 防止除零
    }

    bool try_acquire(uint64_t n = 1) {
        if (n == 0) return true;
        uint64_t now_ms = now();
        size_t idx = (now_ms / slot_size_ms_) % slot_count_;
        uint64_t slot_start_ms = (now_ms / slot_size_ms_) * slot_size_ms_;

        // 尝试初始化当前 slot
        uint64_t prev_start = slots_[idx].start_ms.load(std::memory_order_acquire);
        if (prev_start != slot_start_ms) {
            // CAS 更新起始时间，如果成功则清零计数
            if (slots_[idx].start_ms.compare_exchange_strong(prev_start, slot_start_ms,
                                                              std::memory_order_acq_rel,
                                                              std::memory_order_acquire)) {
                slots_[idx].count.store(0, std::memory_order_relaxed);
            }
        }

        // 原子增加当前 slot
        uint64_t new_count = slots_[idx].count.fetch_add(n, std::memory_order_acq_rel) + n;

        // 计算窗口总请求数
        uint64_t total = 0;
        for (size_t i = 0; i < slot_count_; ++i) {
            total += slots_[i].count.load(std::memory_order_relaxed);
        }

        if (total <= capacity_) {
            return true;
        } else {
            // 超额回退
            slots_[idx].count.fetch_sub(n, std::memory_order_acq_rel);
            return false;
        }
    }

private:
    uint64_t now() const {
        return std::chrono::duration_cast<std::chrono::milliseconds>(
                   clock::now().time_since_epoch())
            .count();
    }

private:
    const uint64_t window_size_ms_;
    const uint64_t capacity_;
    size_t slot_count_;
    uint64_t slot_size_ms_;
    std::vector<Slot> slots_;
};

// ======================= 示例 =======================
int main() {
    MultiSlotSlidingWindowLimiter limiter(1000 /*1s窗口*/, 20 /*容量*/, 10 /*10个slot*/);

    constexpr int kThreads = 4;
    constexpr int kOpsPerThread = 20;
    std::atomic<int> success{0}, fail{0};

    std::vector<std::thread> ths;
    for (int i = 0; i < kThreads; ++i) {
        ths.emplace_back([&]() {
            for (int j = 0; j < kOpsPerThread; ++j) {
                if (limiter.try_acquire()) {
                    success.fetch_add(1, std::memory_order_relaxed);
                } else {
                    fail.fetch_add(1, std::memory_order_relaxed);
                }
                std::this_thread::sleep_for(std::chrono::milliseconds(50));
            }
        });
    }

    for (auto &t : ths) t.join();

    std::cout << "Success=" << success.load() << ", Fail=" << fail.load() << "\n";
    return 0;
}
